# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jhkim

import math
import numpy
from coin.exchange.base.order_gateway import OrderSide


def get_book_range_qty(*, is_buy, book, order_price):
  if is_buy:
    pqarray = numpy.array(book.get_ask_array(20))
    return pqarray[pqarray[:, 0] <= order_price, 1].sum()
  else:
    pqarray = numpy.array(book.get_bid_array(20))
    return pqarray[pqarray[:, 0] >= order_price, 1].sum()


def ceil_price(price, tick_size):
  return math.ceil(price / tick_size) * tick_size


def floor_price(price, tick_size):
  return math.floor(price / tick_size) * tick_size


def get_norm_pos_potential(executor):
  if executor._position is None:
    return 0, 0

  buy_potential_pos = (executor._wo._buy_potential + executor._position - executor.config2.reserve)
  sell_potential_pos = (executor._wo._sell_potential + executor._position
                        - executor.config2.reserve)

  buy_norm_pos = buy_potential_pos / (executor.config.max_pos - executor.config.min_pos) * 2
  sell_norm_pos = sell_potential_pos / (executor.config.max_pos - executor.config.min_pos) * 2

  return buy_norm_pos, sell_norm_pos


class LmAggLogic(object):
  def __init__(self, *, config, trade_product, logger):
    self._config = config
    self._logger = logger
    self._last_manage = 0
    self._threshold_bps = self._config['threshold_bps']
    self._edge_bps = self._config['base_edge_bps']
    self._edge_allowance_bps = self._config['edge_allowance_bps']
    self._model_output_is_bps = self._config['model_output_is_bps']
    self._lean_bps = self._config.get('lean_edge_bps', 0)
    self._price_adj_ratio = self._config.get('price_adj_ratio', 0.0) or 0.0
    self._tick_size = self._config['tick_size']

    assert 0 <= self._price_adj_ratio <= 1.0, self._price_adj_ratio

    self._sfd_max_lot = None
    if trade_product.exchange == "Bitflyer":
      self._sfd_max_lot = 10
      import coin.exchange.bitflyer_v1.order_gateway
      from absl import flags
      flags.FLAGS.bitflyer_quick_ignore = True
      print("using bitflyer quick ignore feature")

  def manage_orders(self, executor, timestamp):
    if executor is None:
      return
    if self._last_manage + 1.5e8 < timestamp:
      executor.manage_agg_orders()
      self._last_manage = timestamp

  def try_fire_and_manage(self, *, timestamp, executor, feed_converter, model, book, dry_run):
    self.try_fire_impl(timestamp=timestamp,
                       executor=executor,
                       feed_converter=feed_converter,
                       model=model,
                       book=book,
                       dry_run=dry_run)
    if not dry_run:
      self.manage_orders(executor, timestamp)

  def try_fire_impl(self, *, timestamp, executor, feed_converter, model, book, dry_run):
    if executor is None:
      return False
    product = executor._product

    ret_proto = feed_converter.trigger(product, model)
    buy_model_output = None
    sell_model_output = None

    ask0_add = 0
    bid0_add = 0
    if ret_proto is not None:
      try:
        ask0_add = ret_proto.ask0_price_adj * self._price_adj_ratio
        bid0_add = ret_proto.bid0_price_adj * self._price_adj_ratio
      except Exception:
        pass
      if len(ret_proto.buy_model) > 0:
        buy_model_output = max([res.signal for res in ret_proto.buy_model])
      if len(ret_proto.sell_model) > 0:
        sell_model_output = max([res.signal for res in ret_proto.sell_model])
      if buy_model_output is None and sell_model_output is None:
        return False
    else:
      return False

    if not (book.has_ask() and book.has_bid()):
      return False

    ask0 = book.ask0().price
    bid0 = book.bid0().price

    prints = []

    prints.append(("%s ask0p_%s bid0p_%s", timestamp, ask0, bid0))

    if abs(ask0_add) > ask0 * 5e-4 or abs(bid0_add) > bid0 * 5e-4:
      prints.append(("fallback to 0: exceed 5bp ask0_add & bid0_add: %s, %s", ask0_add, bid0_add))
      ask0_add = 0
      bid0_add = 0
    else:
      prints.append(("ask0_add: %s, bid0_add: %s", ask0_add, bid0_add))

    mtmp = 0.5 * (ask0 + bid0)
    signal_to_bps_divisor = 1 if self._model_output_is_bps else (mtmp / 1e4)

    buy_model_output_bps = buy_model_output / signal_to_bps_divisor if buy_model_output is not None else None
    sell_model_output_bps = sell_model_output / signal_to_bps_divisor if sell_model_output is not None else None

    buy_on = buy_model_output_bps is not None and buy_model_output_bps > self._threshold_bps
    sell_on = sell_model_output_bps is not None and sell_model_output_bps > self._threshold_bps

    buy_thres = self._threshold_bps
    sell_thres = self._threshold_bps

    if self._lean_bps > 0:
      buy_norm_pos, sell_norm_pos = get_norm_pos_potential(executor)
      if buy_norm_pos > 0:
        buy_thres = self._threshold_bps + self._lean_bps * min(buy_norm_pos, 1.0)
        prints.append(("buy thres changes, %s -> %s", self._threshold_bps, buy_thres))
      if sell_norm_pos < 0:
        sell_thres = self._threshold_bps + self._lean_bps * min(abs(sell_norm_pos), 1.0)
        prints.append(("sell thres changes, %s -> %s", self._threshold_bps, sell_thres))

      if buy_on:
        buy_on = buy_model_output_bps > buy_thres
        prints.append(("thus buy_on: %s", buy_on))
      if sell_on:
        sell_on = sell_model_output_bps > sell_thres
        prints.append(("thus sell_on: %s", sell_on))

    if buy_on and sell_on:
      self._logger.warning('Buy and Sell triggered at once: %s, %s.',
                           buy_model_output_bps,
                           sell_model_output_bps)
      self._logger.warning("No action.")
    elif not buy_on and not sell_on:
      return False

    buy_agg_edge_bps = 0
    sell_agg_edge_bps = 0
    if buy_on:
      buy_agg_edge_bps = buy_model_output_bps - self._edge_bps
      buy_agg_edge_bps = min(buy_agg_edge_bps, self._edge_allowance_bps)
      prints.append(
          ("Bsignal %s -> adj %s, thres %s", buy_model_output_bps, buy_agg_edge_bps, buy_thres))
    if sell_on:
      sell_agg_edge_bps = sell_model_output_bps - self._edge_bps
      sell_agg_edge_bps = min(sell_agg_edge_bps, self._edge_allowance_bps)
      prints.append(
          ("Ssignal %s -> adj %s, thres %s", sell_model_output_bps, sell_agg_edge_bps, sell_thres))

    sell_price = ceil_price(bid0 - sell_agg_edge_bps * (mtmp / 1e4) + bid0_add, self._tick_size)
    buy_price = floor_price(ask0 + buy_agg_edge_bps * (mtmp / 1e4) + ask0_add, self._tick_size)

    if self._sfd_max_lot is not None:
      if buy_on and hasattr(book, "get_ask_price_by_qty"):
        buyp = book.get_ask_price_by_qty(self._sfd_max_lot)
        if buy_price >= buyp:
          prints.append(("buyp adjusted: %s -> %s", buy_price, buyp - self._tick_size))
          buy_price = min(buy_price, buyp - self._tick_size)
      if sell_on and hasattr(book, "get_bid_price_by_qty"):
        sellp = book.get_bid_price_by_qty(self._sfd_max_lot)
        if sell_price <= sellp:
          prints.append(("sellp adjusted: %s -> %s", sell_price, sellp + self._tick_size))
          sell_price = max(sell_price, sellp + self._tick_size)

    res = False

    if not dry_run and executor is not None:
      if buy_on:
        buy_qty = get_book_range_qty(is_buy=True, book=book, order_price=buy_price)
        buy_qty = min(buy_qty, executor.config2.lot_size)
        if buy_qty == 0:
          prints.append(("buy suppressed due to liquidity",))
        else:
          res = executor.submit_agg_order(order_side=OrderSide.BUY,
                                          order_price=buy_price,
                                          order_qty=buy_qty)
          prints.append(("send buy code: %s", executor._prev_buy_code))
      elif sell_on:
        sell_qty = get_book_range_qty(is_buy=False, book=book, order_price=sell_price)
        sell_qty = min(sell_qty, executor.config2.lot_size)
        if sell_qty == 0:
          prints.append(("sell suppressed due to liquidity",))
        else:
          res = executor.submit_agg_order(order_side=OrderSide.SELL,
                                          order_price=sell_price,
                                          order_qty=sell_qty)
          prints.append(("send sell code: %s", executor._prev_sell_code))
    else:
      if buy_on:
        prints.append(("buy fired",))
      elif sell_on:
        prints.append(("sell fired",))

    for printline in prints:
      self._logger.info(*printline)

    return res
